import { Injectable } from '@nestjs/common';
import { ConfigService } from '@nestjs/config';
import { AllConfigType } from '@/config/config.type';
import {
  createAgent,
  createMiddleware,
  // providerStrategy,
  toolStrategy,
  ToolMessage,
  type ReactAgent,
  type tool as Tool,
  type AgentMiddleware,
} from 'langchain';
// import { MessagesZodState } from '@langchain/langgraph';
import * as z from 'zod';
import { ChatOpenAI } from '@langchain/openai';
import { MultiServerMCPClient } from '@langchain/mcp-adapters';
import { PostgresSaver } from '@langchain/langgraph-checkpoint-postgres';
import { LoggerService } from '@/common/logger/logger.service';
import pg from 'pg';
@Injectable()
export class AgentService {
  private agentConfig: AllConfigType['agent'];
  private agent: ReactAgent;
  private tools: Tool[];
  private handleToolErrors: AgentMiddleware;
  private checkpointer: PostgresSaver;
  constructor(
    private configService: ConfigService<AllConfigType>,
    private logger: LoggerService,
  ) {
    this.agentConfig = this.configService.getOrThrow('agent', { infer: true });
    this.init();
  }
  init() {
    Promise.all([this.initShortMemory(), this.initMcp(), this.initErrTool()])
      .then(() => {
        this.initAgent();
      })
      .catch((e) => {
        this.logger.error(e);
      });
  }
  async initMcp() {
    const mcpClient = new MultiServerMCPClient({
      // Global tool configuration options
      // Whether to throw on errors if a tool fails to load (optional, default: true)
      throwOnLoadError: true,
      // Whether to prefix tool names with the server name (optional, default: false)
      prefixToolNameWithServerName: false,
      // Optional additional prefix for tool names (optional, default: "")
      additionalToolNamePrefix: '',

      // Use standardized content block format in tool outputs
      useStandardContentBlocks: true,

      // Server configuration
      mcpServers: {
        // Sreamable HTTP transport example, with auth headers and automatic SSE fallback disabled (defaults to enabled)
        playwright: {
          url: 'http://116.148.216.90:18931/sse',
          headers: {
            Authorization: 'Bearer token123',
          },
        },
      },
    });
    this.tools = await mcpClient.getTools();
  }
  async initShortMemory() {
    const db = this.configService.getOrThrow('database', { infer: true });
    //"postgresql://postgres:postgres@localhost:5442/postgres?sslmode=disable";
    const DB_URI = `postgresql://${db.username}:${db.password}@${db.host}:${db.port}/${db.name}?sslmode=disable`;
    const pool = new pg.Pool({
      connectionString: DB_URI,
    });
    this.checkpointer = new PostgresSaver(pool, undefined, {
      schema: 'agent_memory',
    });
    await this.checkpointer.setup();
  }
  initErrTool() {
    this.handleToolErrors = createMiddleware({
      name: 'HandleToolErrors',
      wrapToolCall: (request, handler) => {
        try {
          return handler(request);
        } catch (error) {
          // Return a custom error message to the model
          return new ToolMessage({
            content: `Tool error: Please check your input and try again. (${error})`,
            tool_call_id: request.toolCall.id!,
          });
        }
      },
    });
  }
  initAgent() {
    const { llmModelName, llmModelApiUrl, llmModelApiKey } = this.agentConfig;
    const systemPrompt = '';
    // 格式化数据
    const responseFormat = z.object({
      name: z.string(),
    });

    const model = new ChatOpenAI({
      model: llmModelName,
      temperature: 0.8,
      maxTokens: 1000,
      timeout: 30000,
      apiKey: llmModelApiKey,
      configuration: {
        baseURL: llmModelApiUrl,
      },
    });
    this.agent = createAgent({
      model: model,
      systemPrompt,
      tools: [...this.tools],
      responseFormat: toolStrategy(responseFormat),
      checkpointer: this.checkpointer,
      middleware: [this.handleToolErrors],
      // logprobs: true,
    }) as unknown as ReactAgent;
  }
  async exec(sessionId: string, content: string) {
    const config = {
      configurable: { thread_id: '1' },
      context: { user_id: sessionId },
    };
    const response = await this.agent.invoke({
      messages: [{ role: 'user', content: content }],
      config,
    });

    return { sessionId, content, response };
  }
}
